# core/storage.py
import sqlite3
import os
import json
from typing import List, Dict, Optional
from core.models import User, Weibo


class UserStorage:
    def __init__(self, conn: sqlite3.Connection):
        self.conn = conn

    def create_table(self):
        cursor = self.conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS author (
                user_id TEXT PRIMARY KEY,
                nickname TEXT,
                gender TEXT,
                location TEXT,
                birthday TEXT,
                description TEXT,
                verified_reason TEXT,
                weibo_count INTEGER,
                following_count INTEGER,
                follower_count INTEGER
            )
        ''')
        self.conn.commit()

    def save_user(self, user: User):
        cursor = self.conn.cursor()
        cursor.execute('''
            INSERT OR IGNORE INTO author VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            user.user_id, user.nickname, user.gender, user.location,
            getattr(user, 'birthday', None), user.description,
            user.verified_reason, user.weibo_count,
            user.following_count, user.follower_count
        ))
        self.conn.commit()


class WeiboStorage:
    def __init__(self, conn: sqlite3.Connection):
        self.conn = conn

    def create_table(self):
        cursor = self.conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS weibo (
                id TEXT PRIMARY KEY,
                user_id TEXT,
                content TEXT,
                publish_time TEXT,
                device TEXT,
                repost_count INTEGER,
                comment_count INTEGER,
                like_count INTEGER,
                images TEXT,
                FOREIGN KEY (user_id) REFERENCES author(user_id)
            )
        ''')
        self.conn.commit()

    def insert_weibo(self, weibo: Dict) -> bool:
        cursor = self.conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM weibo WHERE id=?", (weibo['id'],))
        exists = cursor.fetchone()[0]
        if exists:
            return False

        cursor.execute('''
            INSERT INTO weibo (
                id, user_id, content, publish_time, device, 
                repost_count, comment_count, like_count, images
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            weibo['id'],
            weibo.get('user_id'),
            weibo['content'],
            weibo['publish_time'],
            weibo.get('device', ''),
            weibo.get('repost_count', 0),
            weibo.get('comment_count', 0),
            weibo.get('like_count', 0),
            json.dumps(weibo.get('images', []), ensure_ascii=False)
        ))
        self.conn.commit()
        return True

    def bulk_insert_weibos(self, weibos: List[Dict]) -> int:
        new_count = 0
        for weibo in weibos:
            if self.insert_weibo(weibo):
                new_count += 1
        return new_count

    def count_new_weibos(self, weibos: List[Dict]) -> int:
        ids = tuple(w['id'] for w in weibos)
        if not ids:
            return 0
        query = f"SELECT COUNT(*) FROM weibo WHERE id IN ({','.join(['?']*len(ids))})"
        cursor = self.conn.cursor()
        cursor.execute(query, ids)
        existing_count = cursor.fetchone()[0]
        return len(weibos) - existing_count
